import requests
from PIL import Image
from io import BytesIO
from diffusers import StableDiffusionUpscalePipeline,KDPM2DiscreteScheduler
import torch
from constant import prompt,negative_prompt

# load model and scheduler
model_id = "stabilityai/stable-diffusion-x4-upscaler"
pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id,
                                                          torch_dtype=torch.float16
                                                          )
pipeline = pipeline.to("cuda")
pipeline.enable_xformers_memory_efficient_attention()
pipeline.enable_sequential_cpu_offload()
pipeline.enable_attention_slicing("max")
pipeline.vae.enable_tiling()

low_res_img = Image.open("../../data/original_2023_07_22_12_32_24.png").convert("RGB")


upscaled_image = pipeline(prompt=prompt, image=low_res_img,num_inference_steps=10,negative_prompt=negative_prompt).images[0]
# 当前时间 yyyy-mm-dd-hh-mm-ss
import datetime
now = datetime.datetime.now()
# 转换为指定的格式:
otherStyleTime = now.strftime("%Y_%m_%d_%H_%M_%S")
upscaled_image.save("../data/upsampled_" + otherStyleTime + ".png")
